package leetcode.code1819;

import java.util.PriorityQueue;

public class MKAverage2 {

	int M = (int) 1e5 + 1;
	int lk, rk;
	int m, k;
	int[] where;
	int[] nums;
	int p;
	int sm, div;// 中间累加和
	PriorityQueue<int[]> rkq;// int[]{val,index}
	PriorityQueue<int[]> lkq;
	PriorityQueue<int[]> mbq;
	PriorityQueue<int[]> msq;

	public MKAverage2(int m, int k) {
		this.m = m;
		this.k = k;
		lk = k;
		rk = k;
		where = new int[M];// 0,1,2
		nums = new int[M];// num
		p = 0;// 可用位置
		sm = 0;// 中间和
		div = m - (k << 1);
		rkq = new PriorityQueue<>((a, b) -> a[0] - b[0]);// k大，小根堆
		lkq = new PriorityQueue<>((a, b) -> b[0] - a[0]);// k小，大根堆
		msq = new PriorityQueue<>((a, b) -> a[0] - b[0]);// rest小根堆
		mbq = new PriorityQueue<>((a, b) -> b[0] - a[0]);// rest大根堆
	}

	public void addElement(int num) {
		nums[p] = num;
		available(p - m);// 索引小于p - m都是不可用
		if (rk > 0) {// 大堆缺
			if ((!mbq.isEmpty() && num >= nums[mbq.peek()[1]])) {
				where[p] = 2;// num放在大堆
				rkq.add(new int[] { num, p });
			} else if (!mbq.isEmpty()) {
				int[] c = mbq.poll();
				where[p] = 1;// num放在中间
				where[c[1]] = 2;// 中间移动到大堆
				sm = sm - nums[c[1]] + num;
				rkq.add(new int[] { nums[c[1]], c[1] });
				mbq.add(new int[] { num, p });
				msq.add(new int[] { num, p });
			} else if ((!lkq.isEmpty() && num >= nums[lkq.peek()[1]])) {
				where[p] = 2;// num放在大堆
				rkq.add(new int[] { num, p });
			} else if (!lkq.isEmpty()) {
				int[] c = lkq.poll();
				where[p] = 0;// num放在中间
				where[c[1]] = 2;// 中间移动到大堆
				rkq.add(new int[] { nums[c[1]], c[1] });
				lkq.add(new int[] { num, p });
			} else {
				where[p] = 2;// num放在大堆
				rkq.add(new int[] { num, p });
			}
			rk--;
		} else if (lk > 0) {// 小堆缺
			if ((!msq.isEmpty() && num <= nums[msq.peek()[1]])) {
				where[p] = 0;// num放在小堆
				lkq.add(new int[] { num, p });
			} else if (!msq.isEmpty()) {
				int[] c = msq.poll();
				where[p] = 1;// num放在中间
				where[c[1]] = 0;// 中间移动到大堆
				sm = sm - nums[c[1]] + num;
				lkq.add(new int[] { nums[c[1]], c[1] });
				mbq.add(new int[] { num, p });
				msq.add(new int[] { num, p });
			} else if ((!rkq.isEmpty() && num <= nums[rkq.peek()[1]])) {
				where[p] = 0;// num放在小堆
				lkq.add(new int[] { num, p });
			} else if (!rkq.isEmpty()) {
				int[] c = rkq.poll();
				where[p] = 2;// num
				where[c[1]] = 0;//
				lkq.add(new int[] { nums[c[1]], c[1] });
				rkq.add(new int[] { num, p });
			} else {
				where[p] = 0;// num放在小堆
				lkq.add(new int[] { num, p });
			}
			lk--;
		} else {// 中间缺
			if (num >= nums[rkq.peek()[1]]) {
				int[] rr = rkq.poll();
				int idr = rr[1];
				where[idr] = 1;
				sm = sm + nums[idr];
				msq.add(new int[] { nums[idr], idr });
				mbq.add(new int[] { nums[idr], idr });
				where[p] = 2;
				rkq.add(new int[] { num, p });
			} else if (num <= nums[lkq.peek()[1]]) {
				int[] ll = lkq.poll();
				int idl = ll[1];
				where[idl] = 1;
				sm = sm + nums[idl];
				msq.add(new int[] { nums[idl], idl });
				mbq.add(new int[] { nums[idl], idl });
				where[p] = 0;
				lkq.add(new int[] { num, p });
			} else {
				where[p] = 1;
				sm += num;
				msq.add(new int[] { num, p });
				mbq.add(new int[] { num, p });
			}
		}
		p++;
	}

	// 索引小于p - m都是不可用
	private void available(int i) {
		int pop = i >= 0 ? where[i] : -1;
		if (pop == 0) {// 左端出去的是最小k个里面的数
			lk += 1;
		} else if (pop == 2) {// 左端出去的是最大k个里面的数
			rk += 1;
		} else if (pop == 1) {
			sm -= nums[i];
		}
		// 按照i懒删除一下各个堆，保证堆顶的可用
		poll(rkq, i, 2);
		poll(lkq, i, 0);
		poll(msq, i, 1);
		poll(mbq, i, 1);
	}

	private void poll(PriorityQueue<int[]> q, int i, int type) {
		while (!q.isEmpty() && (q.peek()[1] < i || where[q.peek()[1]] != type)) {// 懒删除
			q.poll();
		}
	}

	public int calculateMKAverage() {
		if (p < m) {
			return -1;
		}
		return sm / div;
	}

	public static void main(String[] args) {

		MKAverage2 obj = new MKAverage2(5, 2);
		obj.addElement(3); // 当前元素为 [3]
		obj.addElement(1); // 当前元素为 [3,1]
		System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
		obj.addElement(10); // 当前元素为 [3,1,10]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
		// 删除最小以及最大的 1 个元素后，容器为 [3]
		// [3] 的平均值等于 3/1 = 3 ，故返回 3
		obj.addElement(5); // 当前元素为 [3,1,10,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]
		obj.addElement(3); // 当前元素为 [3]
		obj.addElement(1); // 当前元素为 [3,1]
		System.out.println(obj.calculateMKAverage()); // 返回 -1 ，因为 m = 3 ，但数据流中只有 2 个元素
		obj.addElement(10); // 当前元素为 [3,1,10]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [3,1,10]
		// 删除最小以及最大的 1 个元素后，容器为 [3]
		// [3] 的平均值等于 3/1 = 3 ，故返回 3
		obj.addElement(5); // 当前元素为 [3,1,10,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5]
		obj.addElement(5); // 当前元素为 [3,1,10,5,5,5]
		System.out.println(obj.calculateMKAverage()); // 最后 3 个元素为 [5,5,5]

	}

}
